import numpy as np
import config
import getGLP
import getBasis
import Error
import plot
import TVDLimter
from matplotlib import pyplot as plt

# 定义参数
# dimPk为分片多项式空间的维数
# dimPk == num_basis == k+1  （一维）
# dimPk 表示张成多项式空间所需的最少基函数个数
# 二维情况下
# dimPk == num_basis == (k+1)(k+2)/2
# P0: 1 P1: 3 P2: 6 P3: 10
Nx = config.Nx  # 单元数
k = config.k  # 分片多项式空间P_k的次数
NumGLP = config.NumGLP  # Gauss点的个数
dimPk = config.dimPk  # 分片多项式空间的维数
CFL = config.CFL  # CFL数


# Gauss点和权重
weight, lambda_ = getGLP.load_GLP(NumGLP)
xa = config.xa  # 区间左端点
xb = config.xb  # 区间右端点
bcL = config.bcL
bcR = config.bcR
t_end = config.t_end  # 终止时间


# init 进行区间变换
t = 0  # 初始时刻
hx = config.hx  # 网格步长
hx1 = config.hx1
Xc = config.Xc  # 单元中心坐标

# dt = CFL * hx / (2*k + 1)
dt = CFL * hx
f = lambda u: u  # 通量函数


def godunov_flux(uL, uR):
    # f(u) = u*u/2
    if uL <= uR:
        if 0 <= uL:     return 0.5*uL*uL
        elif uR <= 0:   return 0.5*uR*uR
        else:           return 0.0
    else:
        return max(0.5*uL*uL, 0.5*uR*uR)
# ...



def Lh(uh: np.ndarray, f=lambda u: u):
    """
    空间离散算子
    输入:
        uh: (Nx, dimPk) 每个单元上的系数/自由度
        f : 通量函数，需支持对 ndarray 的逐元素运算
    输出:
        du: (Nx, dimPk) 离散算子作用后的结果
    """
    du = np.zeros((Nx, dimPk), dtype=float)
    uhG = np.zeros((Nx, NumGLP), dtype=float)
    uhb = np.vstack([np.zeros((1, dimPk), dtype=float), uh, np.zeros((1, dimPk), dtype=float)])
    uhR = np.zeros((Nx+1, 1), dtype=float)
    uhL = np.zeros((Nx+1, 1), dtype=float)
    f_hat = np.zeros((Nx+1, 1), dtype=float)
    
    # 周期边界 1D情况: u(0, t) = u(L, t)
    # 数值通量需要两侧单元的解来进行计算 
    if bcL == 1:
        uhb[0, :] = uh[-1, :]    # 左幽灵 = 最后一个物理单元
    elif bcL == 2:
        uhb[0, :] = uh[1, :]     # 左幽灵 = 第一个物理单元
    if bcR == 1:
        uhb[-1, :] = uh[0, :]    # 右幽灵 = 第一个物理单元
    elif bcR == 2:
        uhb[-1, :] = uh[-2, :]   # 右幽灵 = 最后一个物理单元

    # --- Step 1: 单元内部高斯点插值 uhG 与体积分 ---
    for i in range(Nx):
        for d in range(dimPk):
            # uhG[i,d] = uhb[i,:]·phiG[i,d] (来自单元 i 的 uhb 的线性组合)
            uhG[i] +=  uh[i, d] * getBasis.phiG[:, d].T

    for i in range(Nx):
        for d in range(1, dimPk):
            for i1 in range(NumGLP):
                du[i, d] += 0.5 * weight[i1] * f(uhG[i, i1]) * getBasis.phixG[i1, d]

    # --- Step 2: 边界数值通量 ---
    for i in range(Nx+1):
        for d in range(dimPk):
            uhR[i] += uhb[i, d] * getBasis.phiGR[d]
            uhL[i] += uhb[i+1, d] * getBasis.phiGL[d]
            
    for i in range(Nx+1):
        uR = uhL[i]
        uL = uhR[i]
        # alpha = max(|f'(u)|)
        alpha = max(abs(uL), abs(uR))  # 这里使用Lax-Friedrichs通量
        f_hat[i] = 0.5 * (f(uR) + f(uL) - alpha * (uR - uL))
        # f_hat[i] = godunov_flux(float(uL), float(uR))

    # --- Step 3: 面项贡献到每个单元 ---
    for i in range(Nx):
        for d in range(dimPk):
            du[i, d] += -1/hx * (getBasis.phiGR[d] * f_hat[i+1] - getBasis.phiGL[d] * f_hat[i])

    # --- Step 4: 质量矩阵对角项归一化 ---
    for d in range(dimPk):
        du[:, d] /= getBasis.mm[d]

    return du
    

u_real = np.zeros((Nx, NumGLP))  # 储存初值在每个高斯点上的值
for i in range(Nx):
    for j in range(NumGLP):
        x = hx1 * lambda_[j] + Xc[i]  # 变换到物理区间
        u_real[i, j] = config.u0(x)  # 真实解

weight, lambda_ = getGLP.load_GLP(config.NumGLP)
uh = np.zeros((config.Nx, config.dimPk))  # 储存每个单元的系数

# L2 Projection  将连续初始场u0投影到离散空间
# \int_{\Omega} (u_real * phi(x)) dx
# 在高斯点上的数值，逐点乘上权重
for i in range(Nx):
    for d in range(dimPk):
        for i1 in range(NumGLP):
            uh[i, d] += 0.5 * weight[i1] * u_real[i, i1] * getBasis.phiG[i1, d]
    
for d in range(dimPk):
    uh[:, d] /= getBasis.mm[d]


# RK3 计算的核心部分
while t < t_end:
    if t + dt >= t_end:
        dt = t_end - t
        t = t_end
    else:
        t += dt
    
    # RK3公式
    du = Lh(uh, f=f)
    uh1 = uh + dt * du
    uh1 = TVDLimter.TVD_Limiter(uh1, Nx, hx, dimPk, bcL, bcR, M=config.M)  # TVD 限制器
    
    
    # stage II
    du = Lh(uh1, f=f)
    uh2 = 3.0 / 4.0 * uh + 1.0 / 4.0 * (uh1 + dt * du)
    uh2 = TVDLimter.TVD_Limiter(uh2, Nx, hx, dimPk, bcL, bcR, M=config.M)  # TVD 限制器
    
    # stage III
    du = Lh(uh2, f=f)
    uh = 1.0 / 3.0 * uh + 2.0 / 3.0 * (uh2 + dt * du)
    uh = TVDLimter.TVD_Limiter(uh, Nx, hx, dimPk, bcL, bcR, M=config.M)  # TVD 限制器
    
plot.plot_solution(uh, t, weight, lambda_, Xc, hx1, title_prefix="DG solution", filename="figures/rarefaction.png")


u_exact_tend = np.zeros((Nx, NumGLP))  # t_end时刻的解析解
for i in range(Nx):
    for j in range(NumGLP):
        x = hx1 * lambda_[j] + Xc[i]
        u_exact_tend[i, j] = config.true_u(x)  # 注意：true_u 必须写对（见下一节）
    
L2_err = Error.calculate_L2_Error(
    uh=uh,
    u_real=u_exact_tend,
    Nx=Nx,
    NumGLP=NumGLP,
    dimPk=dimPk,
    phiG=getBasis.phiG,
    weight=weight,
    hx1=hx1
)

# uh: (Nx, dimPk)  存储了每个单元的系数
# 利用
print(f"Error = {L2_err}")